﻿// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved. 
// See License.txt in the project root for license information.
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.Contracts;
using System.Globalization;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;

namespace System.Threading.Actors.Utilities
{
    // Common assembly-related logic that is used in more than one location.
    public class AssemblyHelper
    {
        // An AppDomain-wide cache of assemblies, to enable:
        //      (1) Referencing assemblies in actor state, instead of storing entire assemblies
        //      (2) Supporting the storage/referencing of multiple same-named assemblies, distinguished by checksum value
        //
        // The intent here is that the key is an assembly name concatenated with a stringified checksum.
        // This will allow same-named assemblies to be tracked in the cache.  The disadvantage is that
        // outdated assemblies will never be evicted from the cache.
        internal static readonly Dictionary<string, byte[]> s_assemblyCache = new Dictionary<string, byte[]>();

        private IActorState m_state; // Need this for general access to state
        private ConcurrentDictionary<string, object> m_permaStore; // Need this to store/remove assembly/method information at the lowest level
        private MethodCallHelper m_methodCallHelper;

        public AssemblyHelper(IActorState state, ConcurrentDictionary<string,object> permaStore, MethodCallHelper methodCallHelper)
        {
            Contract.Requires(state != null, "Expected state to be non-null");
            Contract.Requires(permaStore != null, "Expected permaStore to be non-null");
            Contract.Requires(methodCallHelper != null, "Expected methodCallHelper to be non-null");
            
            m_state = state;
            m_permaStore = permaStore;
            m_methodCallHelper = methodCallHelper;
        }

        // Scours assembly for actor methods
        private void ScourAssembly(Assembly assembly, bool addMethods = true)
        {
            Contract.Requires(assembly != null);

            var assemblyName = assembly.GetName().Name;
            var scouredAssemblyName = assembly.ManifestModule.ScopeName;
            List<string> methodsAdded = new List<string>(); // Tracks methods that we've loaded
            List<string> methodsRemoved = new List<string>(); // Tracks methods that we've unloaded

            foreach (var module in assembly.GetModules())
            {
                // Got this trick from
                //  http://haacked.com/archive/2012/07/23/get-all-types-in-an-assembly.aspx
                // It allows you to proceed even though some types might not get loaded correctly.
                IEnumerable<Type> usableTypes;
                try
                {
                    usableTypes = module.GetTypes();
                }
                catch (ReflectionTypeLoadException e)
                {
                    //Console.WriteLine("CAUTION: LoadAssembly had a failure: {0}", e.LoaderExceptions[0].Message);
                    usableTypes = e.Types.Where(t => t != null);
                    foreach (var ex in e.LoaderExceptions) ActorETWEventSource.Log.AssemblyLoadFailure(m_state.GetId(), m_state.PartitionName, scouredAssemblyName, ex.Message);
                }
                foreach (var type in usableTypes)
                {
                    var methods = type.GetMethods(BindingFlags.Public | BindingFlags.Static);
                    foreach (var method in methods)
                    {
                        // TODO: Check that method has the correct signature
                        foreach (var attr in method.GetCustomAttributes(true))
                        {
                            if (attr is ActorMethodAttribute)
                            {
                                if (addMethods)
                                {
                                    var firstParam = method.GetParameters()[0];
                                    var methodName = method.Name;

                                    if (firstParam.ParameterType == typeof(IActorState))
                                    {
                                        try
                                        {
                                            var func = (Func<IActorState, object[], object>)Delegate.CreateDelegate(typeof(Func<IActorState, object[], object>), method);
                                            m_permaStore[method.Name] = func; // Store directly into durable store
                                            methodsAdded.Add(methodName);
                                            //this.Logger.WriteLine("", "ScourAssembly: Added method {0} from assembly {1}", methodName, assemblyName);
                                        }
                                        catch (Exception)
                                        {
                                            ActorETWEventSource.Log.ActorMethodSignatureFailure(m_state.GetId(), m_state.PartitionName, assemblyName, methodName);
                                        }
                                    }
                                    else if (firstParam.ParameterType.IsAssignableFrom(typeof(System.Dynamic.DynamicObject)))
                                    {
                                        try
                                        {
                                            // dynamic object
                                            var func = (Func<dynamic, dynamic[], object>)Delegate.CreateDelegate(typeof(Func<dynamic, dynamic[], object>), method);
                                            m_permaStore[method.Name] = func; // Store directly into durable store
                                            methodsAdded.Add(methodName);
                                            //this.Logger.WriteLine("", "ScourAssembly: Added dynamic state method {0} from assembly {1}", methodName, assemblyName);
                                        }
                                        catch (Exception)
                                        {
                                            ActorETWEventSource.Log.ActorMethodSignatureFailure(m_state.GetId(), m_state.PartitionName, assemblyName, methodName);
                                        }
                                    }
                                    else
                                    {
                                        ActorETWEventSource.Log.ActorMethodSignatureFailure(m_state.GetId(), m_state.PartitionName, assemblyName, methodName);
                                    }
                                }
                                else
                                {
                                    // We need to remove the methods associated with the assembly
                                    object temp;
                                    if (m_permaStore.TryGetValue(method.Name, out temp))
                                    {
                                        // Check to see whether this is a method with dynamic actor state
                                        var dynamicMethodDelegate = temp as Func<dynamic, dynamic[], object>;
                                        // Check that this method actually comes from this assembly;
                                        // it may have been replaced by another assembly.
                                        if (dynamicMethodDelegate != null && dynamicMethodDelegate.Method.DeclaringType.Assembly.ManifestModule.ScopeName == assembly.ManifestModule.ScopeName)
                                        {
                                            var methodAssemblyName = dynamicMethodDelegate.Method.DeclaringType.Assembly.ManifestModule.ScopeName;
                                            if (methodAssemblyName.Equals(scouredAssemblyName))
                                            {
                                                m_permaStore.TryRemove(method.Name, out temp);
                                                methodsRemoved.Add(method.Name);
                                                //this.Logger.WriteLine("", "ScourAssembly: Removed dynamic state method {0} from expiring assembly {1}", method.Name, scouredAssemblyName);
                                            }
                                        }
                                        else
                                        {
                                            // Must be a method with IActorState actor state
                                            var methodDelegate = temp as Func<IActorState, object[], object>;
                                            // Check that this method actually comes from this assembly;
                                            // it may have been replaced by another assembly.
                                            if (methodDelegate != null && methodDelegate.Method.DeclaringType.Assembly.ManifestModule.ScopeName == assembly.ManifestModule.ScopeName)
                                            {
                                                var methodAssemblyName = methodDelegate.Method.DeclaringType.Assembly.ManifestModule.ScopeName;
                                                if (methodAssemblyName.Equals(scouredAssemblyName))
                                                {
                                                    m_permaStore.TryRemove(method.Name, out temp);
                                                    methodsRemoved.Add(method.Name);
                                                    //this.Logger.WriteLine("", "ScourAssembly: Removed method {0} from expiring assembly {1}", method.Name, scouredAssemblyName);
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            if (addMethods)
            {
                if (methodsAdded.Count == 0)
                {
                    //this.Logger.WriteLine("", "ScourAssembly: Added an assembly to the actor but it didn't contain actor methods.  Assembly: {0}", assembly.FullName);
                    ActorETWEventSource.Log.AssemblyHadNoActorMethods(m_state.GetId(), m_state.PartitionName, assembly.FullName);
                }
                else
                {
                    ActorETWEventSource.Log.ActorMethodsLoaded(m_state.GetId(), m_state.PartitionName, assemblyName, methodsAdded);
                }
            }
            else if (methodsRemoved.Count > 0)
            {
                ActorETWEventSource.Log.ActorMethodsUnloaded(m_state.GetId(), m_state.PartitionName, assemblyName, methodsRemoved);
            }
        }

        // Find the stale version of an assembly that is in the process of being loaded.
        // The intent here is to identify the assembly being replaced, so that we
        // can unload methods from that assembly.
        private Assembly FindOldAssembly(string assemblyKeyName)
        {
            // The assembly name could be full or partial (the Json serializer uses partial names),
            // but GetAssemblyKey will handle either.
            String refKeyName = "$REF" + assemblyKeyName; // The key under which this assembly would be stored as an assembly reference

            Object foundValue;
            byte[] assemblyBytes;

            // Look up the assembly.
            // We will be probing the dictionary directly, as we are interested in locating the stale assembly,
            // not the new assembly which is in the process of being replicated.
            if (m_permaStore.TryGetValue(assemblyKeyName, out foundValue))
            {
                // This actor has the assembly stored.
                var tuple = foundValue as Tuple<string, byte[]>;
                assemblyBytes = tuple.Item2;
                Assembly a = Assembly.Load(assemblyBytes);
                return a;
            }
            else if (m_permaStore.TryGetValue(refKeyName, out foundValue))
            {
                // This actor has a reference to the assembly stored
                var chksum = foundValue as string;
                var cacheKey = assemblyKeyName + chksum;
                byte[] cachedAssemblyBytes = null;
                lock (s_assemblyCache)
                {
                    s_assemblyCache.TryGetValue(cacheKey, out cachedAssemblyBytes);
                }
                if (cachedAssemblyBytes != null)
                    return Assembly.Load(cachedAssemblyBytes);
            }

            return null;
        }

        // Returns true if we handled the storing of the key-value pair from this method
        public bool CheckForAssemblyOrMethod(string key, object value, bool scour)
        {
            Contract.Requires(key != null);

            // TODO: Name $ASSEMBLY$ and language pack method keys in such a way that we
            // can do a quick exit here.  For example, maybe the assembly key prefix is $$_ASSEMBLY
            // and language pack method key prefixes are of the form $$_js:, so that a quick check
            // that key didn't begin with $$_ could allow us to quickly return from this method.

            if (key.StartsWith("$ASSEMBLY$"))
            {
                string finalKey;
                object finalValue;

                // Can we convert this into an assembly reference?
                var tuple = value as Tuple<string, byte[]>;
                var cacheKeyName = key + tuple.Item1;
                byte[] cachedAssemblyBytes;
                Assembly assembly;
                object temp;

                bool assemblyIsCached;
                lock (s_assemblyCache)
                {
                    assemblyIsCached = s_assemblyCache.TryGetValue(cacheKeyName, out cachedAssemblyBytes);
                }

                if (assemblyIsCached)
                {
                    // Yes, we have this assembly cached.
                    finalKey = "$REF" + key;
                    finalValue = tuple.Item1; // the assembly's checksum
                    assembly = Assembly.Load(cachedAssemblyBytes);
                    // If we were holding an old version of this assembly, then
                    // get rid of it.
                    m_permaStore.TryRemove(key, out temp);
                }
                else
                {
                    // No cached assembly available.  Store as fully inflated assembly.
                    finalKey = key;
                    finalValue = value;
                    assembly = Assembly.Load(tuple.Item2);
                    lock (s_assemblyCache)
                    {
                        s_assemblyCache[cacheKeyName] = tuple.Item2; // Add/replace the entry in the assembly cache
                    }
                }

                var oldAssembly = FindOldAssembly(key);
                if (scour && (oldAssembly != null)) ScourAssembly(oldAssembly, addMethods: false); // One round to remove existing methods from old assembly
                if (scour && (assembly != null)) ScourAssembly(assembly, addMethods: true); // One round to add new methods.

                m_permaStore[finalKey] = finalValue;
                return true;
            }
            else // Check to see whether this key needs to be handled as a method
            {
                // Check language packs for any recognition of this key
                foreach (var kvp in m_methodCallHelper.LanguagePacks)
                {
                    var pack = kvp.Value;
                    if (pack.OwnsKey(key))
                    {
                        pack.AddMethod((string)value); // just re-adds the implementation; key already exists
                        break;
                    }
                }
            }

            return false;

        }

        public void AddAssemblyResolver()
        {
            AppDomain.CurrentDomain.AssemblyResolve += ActorAssemblyResolver;
        }

        public void RemoveAssemblyResolver()
        {
            AppDomain.CurrentDomain.AssemblyResolve -= ActorAssemblyResolver;
        }

        // Converts an assembly name (either full or partial) to a key useful for looking assemblies up
        // in actor state, as our own replicated GAC.
        internal static string GetAssemblyKey(string assemblyName/*, bool includeVersion = true*/)
        {
            string key;
            if (assemblyName.Contains("PublicKeyToken"))
            {
                AssemblyName name = new AssemblyName(assemblyName);
                key = String.Format("$ASSEMBLY${0}", name.Name);
            }
            else
            {
                key = String.Format("$ASSEMBLY${0}", assemblyName);
            }

            // By convention, let's strip off the suffix.
            if (key.EndsWith(".dll", ignoreCase: true, culture: CultureInfo.InvariantCulture) ||
               key.EndsWith(".exe", ignoreCase: true, culture: CultureInfo.InvariantCulture))
            {
                key = key.Substring(0, key.Length - 4);
            }

            return key;
        }

        // We have a significant challenge with this method.  This is appdomain-wide, but an individual actor
        // may have only certain versions of an assembly added to it, or multiple actors could have
        // multiple assemblies with the same name (ie, two different Test.exe's).  Since this method kicks
        // in for an entire appdomain, we would ideally like to scope its effectiveness to precisely
        // one actor.  We need to understand how significant a problem this is.  However, here's what a 
        // solution might look like.  In every place we run an actor method (ie, IActor's CallMethod and any
        // place we might do pub/sub responses), we add a marker (such as the current IActorState) to the
        // thread's execution context.  In this method, we then check to see if the current execution context
        // contains the same marker.  This is assuming that the CLR will call the assembly resolve event with
        // the original thread's execution context, which may be a bit of a leap.  
        // Before building that solution, we need to ensure that there isn't a better way to handle this.
        // IE, do we use this method significantly, or only for custom user defined data types?
        // If it's the latter, perhaps moving to JsonObjects would allow us to not use this method.
        // This method isn't necessary for running tests involving CloudList<int>, for example.
        private Assembly ActorAssemblyResolver(object sender, ResolveEventArgs args)
        {
            Contract.Assert(m_state != null, "m_state is null in ActorServiceHelper.ActorAssemblyResolver.  Why?");

            // Reject these before taking lock.  An IActorState.GetProxy call, from within the state-lock, could:
            //   --Construct a WrappedTcpClient on a separate thread (but wait for that thread to complete), which could
            //   --Cause the WrappedTcpClient static constructor to run, which could
            //   --Cause the LightupITcpClientFactory static constructor to run, which could
            //   --Attempt to resolve a non-existent type (a type that doesn't exist in the currently loaded assemblies), which would
            //   --Cause this method to be called.
            // In such a scenario, we don't want to attempt to take the lock below, because that would result in deadlock.
            // So, detect and avoid that scenario.
            if (LightupITcpClientFactory.s_ClientProviderCandidates.ContainsKey(args.Name))
            {
                return null;
            }

            // The assembly name could be full or partial (the Json serializer uses partial names),
            // but GetAssemblyKey will handle either.
            String assemblyKeyName = GetAssemblyKey(args.Name/*, true*/);  // The key under which this assembly would be stored as an assembly
            String refKeyName = "$REF" + assemblyKeyName; // The key under which this assembly would be stored as an assembly reference

            Object foundValue;
            byte[] assemblyBytes;

            try
            {
                // HACK HACK HACK
                // This method is called by CallMethod after CallMethod already has a lock. 
                // While using Monitor.Enter worked reentrantly, SemaphoreSlim obviously doesn't.
                // For the time being, let's try without a lock.
                //m_state.AsyncStateLock.Wait();
                //lock (m_state)
                //{
                try
                {
                    // Look up the assembly.
                    if (m_state.TryGet(assemblyKeyName, out foundValue))
                    {
                        // This actor has the assembly stored.
                        var tuple = foundValue as Tuple<string, byte[]>;
                        assemblyBytes = tuple.Item2;
                        Assembly a = Assembly.Load(assemblyBytes);
                        //this.Logger.WriteLine("ReplicatingActorServiceBase", "Loaded assembly {0} from actor state.", a.FullName);
                        return a;
                    }
                    else if (m_state.TryGet(refKeyName, out foundValue))
                    {
                        // This actor has a reference to the assembly stored
                        var chksum = foundValue as string;
                        var cacheKey = assemblyKeyName + chksum;
                        byte[] cachedAssemblyBytes = null;
                        lock (s_assemblyCache)
                        {
                            s_assemblyCache.TryGetValue(cacheKey, out cachedAssemblyBytes);
                        }
                        if (cachedAssemblyBytes != null)
                        {
                            var cachedAssembly = Assembly.Load(cachedAssemblyBytes);
                            //this.Logger.WriteLine("ReplicatingActorServiceBase", "Loaded assembly {0} from assembly cache.", cachedAssembly.FullName);
                            return cachedAssembly;
                        }
                    }
                }
                catch (ObjectDisposedException)
                {
                    // There is a possibility of a race between closing (perhaps load balancing) a service and this 
                    // static, appdomain-wide assembly resolver.  We can safely ignore this.
                    return null;
                }
            }
            finally {
                // HACK HACK HACK - not releasing the lock because we're not taking it.  See above.
                //m_state.AsyncStateLock.Release();
            }

            //this.Logger.WriteLine("ReplicatingActorServiceBase", "Request to loaded assembly {0} from actor state failed.  No assembly.", args.Name);
            return null;
        }

    }
}
